

#include "Coroutine.h"

#include <atomic>
#include <cassert>
#include <iostream>
#include <pthread.h>


static bool debug = false;
namespace zepco {

// 当前线程正在运行的协程
static thread_local Zephyr *t_zephyr = nullptr;
// 当前线程的主协程
static thread_local std::shared_ptr<Zephyr> t_thread_zephyr = nullptr;
// 当前线程的调度协程
static thread_local Zephyr *t_scheduler_zephyr = nullptr;
// 协程ID计数器
static std::atomic<uint64_t> s_zephyr_id{0};
// 活跃协程计数器
static std:: atomic<uint64_t> s_zephyr_count{0};

void Zephyr::SetThis(Zephyr *z) {
    t_zephyr = z;
}

std::shared_ptr<Zephyr> Zephyr::GetThis() {
    if (t_zephyr) {
        return t_zephyr->shared_from_this();
    }

    std::shared_ptr<Zephyr> main_zephyr(new Zephyr());
    t_thread_zephyr = main_zephyr;
    t_scheduler_zephyr = main_zephyr.get(); // 设置调度协程 默认为主协程
    assert(t_zephyr == main_zephyr.get()); // 判断当前运行协程是否为主协程
    return t_zephyr->shared_from_this();
}

void Zephyr::SetSchedulerZephyr(Zephyr *z) {
    t_scheduler_zephyr = z;
}

uint64_t Zephyr::GetZephyrID() {
    if (t_zephyr)
        return t_zephyr->getID();
    return (uint64_t) -1;
}

Zephyr::Zephyr() {
    SetThis(this);
    m_state = RUNNING;

    if (getcontext(&m_ctx)) {
        std::cerr << "Zephyr() failed.\n";
        pthread_exit(nullptr);
    }
    m_id = s_zephyr_id ++;
    s_zephyr_count ++;
    if (debug)
        std::cout << "Zephyr(): main id: " << m_id << std::endl;
}

Zephyr::Zephyr(std::function<void()> callback, uint32_t stacksize, bool run_in_scheduler)
    : m_callback(callback), m_runInScheduler(run_in_scheduler) {
    m_state = READY;
    m_stacksize = stacksize ? stacksize : 128000;
    m_stack = malloc(m_stacksize);
    if (getcontext(&m_ctx)) {
        std::cerr << "Zephyr(std::function<void()> callback, uint32_t stacksize, bool run_in_scheduler) faild\n";
        pthread_exit(nullptr);
    }
    m_ctx.uc_stack.ss_sp = m_stack;
    m_ctx.uc_stack.ss_size = m_stacksize;
    m_ctx.uc_link = nullptr; // ?? MainFunc 结束后会调用yield返回主协程或调度协程
    makecontext(&m_ctx, &Zephyr::MainFunc, 0);

    m_id = s_zephyr_id ++;
    s_zephyr_count ++;
    if (debug)
        std::cout << "Zephyr(): child id: " << m_id << std::endl;
}

Zephyr::~Zephyr() {
    s_zephyr_count --;
    if (m_stack) {
        // 有栈 说明是子协程
        assert(m_state == TERM);
        free(m_stack);
    } else {
        // 无栈 说明是主协程
        assert(m_state == RUNNING);
        Zephyr *current = t_zephyr;
        if (current == this) {
            SetThis(nullptr);
        }
    }

    if(debug)
        std::cout << "~Zephyr(): id = " << m_id << std::endl;
}

void Zephyr::reset(std::function<void()> callback) {
    // 重用已经终止的协程，避免创建销毁带来的开销
    assert(m_stack != nullptr && m_state == TERM);
    m_state = READY;
    m_callback = callback;
    if (getcontext(&m_ctx)) {
        std::cerr << "reset() failed\n";
        pthread_exit(nullptr);
    }

    m_ctx.uc_stack.ss_sp = m_stack;
    m_ctx.uc_stack.ss_size = m_stacksize;
    m_ctx.uc_link = nullptr;
    makecontext(&m_ctx, &Zephyr::MainFunc, 0);
}

void Zephyr::resume() {
    assert(m_state == READY);
    m_state = RUNNING;
    if (m_runInScheduler) {
        if (debug)
            std::cout << t_scheduler_zephyr->m_id << " resume " << m_id << std::endl;
        SetThis(this); // 设置目前工作协程为this
        if (swapcontext(&(t_scheduler_zephyr->m_ctx), &m_ctx)) {
            std::cerr << "resume() from t_scheduler_zephyr failed\n";
            pthread_exit(nullptr);
        }
    } else {
        if (debug)
            std::cout << t_thread_zephyr->m_id << " resume " << m_id << std::endl;
        SetThis(this); // 设置目前工作协程为this
        if (swapcontext(&(t_thread_zephyr->m_ctx), &m_ctx)) {
            std::cerr << "resume() from t_thread_zephyr failed\n";
            pthread_exit(nullptr);
        }
    }
}

void Zephyr::yield() {
    assert(m_state == RUNNING || m_state == TERM);
    if (m_state != TERM)
        m_state = READY;

    if (m_runInScheduler) {
        if (debug)
            std::cout << m_id << " yield " << t_scheduler_zephyr->m_id << std::endl;
        SetThis(t_scheduler_zephyr);
        if (swapcontext(&m_ctx, &(t_scheduler_zephyr->m_ctx))) {
            std::cerr << "yield() to t_scheduler_zephyr faild\n";
            pthread_exit(nullptr);
        }
    } else {
        if (debug)
            std::cout << m_id << " yield " << t_thread_zephyr->m_id << std::endl;
        SetThis(t_thread_zephyr.get());
        if (swapcontext(&m_ctx, &(t_thread_zephyr->m_ctx))) {
            std::cerr << "yield() to t_scheduler_zephyr faild\n";
            pthread_exit(nullptr);
        }
    }
}

void Zephyr::MainFunc() {
    auto current = GetThis();
    assert(current != nullptr);

    current->m_callback();

    current->m_callback = nullptr;
    current->m_state = TERM;

    // 协程函数运行结束时协程即结束，协程结束时会自动调用一次yield以返回主协程
    auto raw_ptr = current.get(); // ?? TODO:
    current.reset(); // 引用计数减一
    raw_ptr->yield(); // ?? TODO:
}

}
